from pydantic import BaseModel, Field, model_validator
from typing import Literal, Self


class Turn(BaseModel):
    speaker: str = Field(default="")
    dia_id: str = Field(default="")
    text: str = Field(default="")
    blip_caption: str = Field(default="")


class Session(BaseModel):
    datetime: str = Field(default="")
    turns: list[Turn] = Field(default_factory=list)


class QA(BaseModel):
    question: str = Field(default="")
    answer: str = Field(default="")
    evidence: list[str] = Field(default_factory=list)
    category: Literal[1, 2, 3, 4, 5] = Field(default="")
    adversarial_answer: str = Field(default="")
    final_answer: str = Field(default="")

    @model_validator(mode="after")
    def get_final_answer(self) -> Self:
        if self.category == 5:
            self.final_answer = self.adversarial_answer
        else:
            self.final_answer = self.answer
        return self
